{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 32、基于PyTorch的文本分类项目模型与训练代码讲解"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# torchtext 不再维护，无法运行了"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "ename": "OSError",
     "evalue": "/home/jcheng/.conda/envs/dl/bin/../lib/libstdc++.so.6: version `GLIBCXX_3.4.30' not found (required by /home/jcheng/.conda/envs/dl/lib/python3.11/site-packages/torchtext/lib/libtorchtext.so)",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mOSError\u001b[39m                                   Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 1\u001b[39m\n\u001b[32m----> \u001b[39m\u001b[32m1\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorchtext\u001b[39;00m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/site-packages/torchtext/__init__.py:18\u001b[39m\n\u001b[32m     15\u001b[39m     _WARN = \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[32m     17\u001b[39m \u001b[38;5;66;03m# the following import has to happen first in order to load the torchtext C++ library\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m18\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorchtext\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m _extension  \u001b[38;5;66;03m# noqa: F401\u001b[39;00m\n\u001b[32m     20\u001b[39m _TEXT_BUCKET = \u001b[33m\"\u001b[39m\u001b[33mhttps://download.pytorch.org/models/text/\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m     22\u001b[39m _CACHE_DIR = os.path.expanduser(os.path.join(_get_torch_home(), \u001b[33m\"\u001b[39m\u001b[33mtext\u001b[39m\u001b[33m\"\u001b[39m))\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/site-packages/torchtext/_extension.py:64\u001b[39m\n\u001b[32m     59\u001b[39m     \u001b[38;5;66;03m# This import is for initializing the methods registered via PyBind11\u001b[39;00m\n\u001b[32m     60\u001b[39m     \u001b[38;5;66;03m# This has to happen after the base library is loaded\u001b[39;00m\n\u001b[32m     61\u001b[39m     \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorchtext\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m _torchtext  \u001b[38;5;66;03m# noqa\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m64\u001b[39m \u001b[43m_init_extension\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/site-packages/torchtext/_extension.py:58\u001b[39m, in \u001b[36m_init_extension\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m     55\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m _mod_utils.is_module_available(\u001b[33m\"\u001b[39m\u001b[33mtorchtext._torchtext\u001b[39m\u001b[33m\"\u001b[39m):\n\u001b[32m     56\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m(\u001b[33m\"\u001b[39m\u001b[33mtorchtext C++ Extension is not found.\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m---> \u001b[39m\u001b[32m58\u001b[39m \u001b[43m_load_lib\u001b[49m\u001b[43m(\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mlibtorchtext\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[32m     59\u001b[39m \u001b[38;5;66;03m# This import is for initializing the methods registered via PyBind11\u001b[39;00m\n\u001b[32m     60\u001b[39m \u001b[38;5;66;03m# This has to happen after the base library is loaded\u001b[39;00m\n\u001b[32m     61\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorchtext\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m _torchtext\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/site-packages/torchtext/_extension.py:50\u001b[39m, in \u001b[36m_load_lib\u001b[39m\u001b[34m(lib)\u001b[39m\n\u001b[32m     48\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m path.exists():\n\u001b[32m     49\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m50\u001b[39m \u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mops\u001b[49m\u001b[43m.\u001b[49m\u001b[43mload_library\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     51\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/site-packages/torch/_ops.py:1392\u001b[39m, in \u001b[36m_Ops.load_library\u001b[39m\u001b[34m(self, path)\u001b[39m\n\u001b[32m   1387\u001b[39m path = _utils_internal.resolve_library_path(path)\n\u001b[32m   1388\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m dl_open_guard():\n\u001b[32m   1389\u001b[39m     \u001b[38;5;66;03m# Import the shared library into the process, thus running its\u001b[39;00m\n\u001b[32m   1390\u001b[39m     \u001b[38;5;66;03m# static (global) initialization code in order to register custom\u001b[39;00m\n\u001b[32m   1391\u001b[39m     \u001b[38;5;66;03m# operators with the JIT.\u001b[39;00m\n\u001b[32m-> \u001b[39m\u001b[32m1392\u001b[39m     \u001b[43mctypes\u001b[49m\u001b[43m.\u001b[49m\u001b[43mCDLL\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   1393\u001b[39m \u001b[38;5;28mself\u001b[39m.loaded_libraries.add(path)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/ctypes/__init__.py:376\u001b[39m, in \u001b[36mCDLL.__init__\u001b[39m\u001b[34m(self, name, mode, handle, use_errno, use_last_error, winmode)\u001b[39m\n\u001b[32m    373\u001b[39m \u001b[38;5;28mself\u001b[39m._FuncPtr = _FuncPtr\n\u001b[32m    375\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m handle \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m376\u001b[39m     \u001b[38;5;28mself\u001b[39m._handle = \u001b[43m_dlopen\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmode\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    377\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m    378\u001b[39m     \u001b[38;5;28mself\u001b[39m._handle = handle\n",
      "\u001b[31mOSError\u001b[39m: /home/jcheng/.conda/envs/dl/bin/../lib/libstdc++.so.6: version `GLIBCXX_3.4.30' not found (required by /home/jcheng/.conda/envs/dl/lib/python3.11/site-packages/torchtext/lib/libtorchtext.so)"
     ]
    }
   ],
   "source": [
    "import torchtext"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "ename": "OSError",
     "evalue": "/home/jcheng/.conda/envs/dl/lib/python3.11/site-packages/torchtext/lib/libtorchtext.so: undefined symbol: _ZN5torch6detail10class_baseC2ERKSsS3_SsRKSt9type_infoS6_",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mOSError\u001b[39m                                   Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[1]\u001b[39m\u001b[32m, line 4\u001b[39m\n\u001b[32m      2\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorch\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mnn\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mnn\u001b[39;00m\n\u001b[32m      3\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorch\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mnn\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mfunctional\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mas\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mF\u001b[39;00m\n\u001b[32m----> \u001b[39m\u001b[32m4\u001b[39m \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorchtext\u001b[39;00m\n\u001b[32m      5\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorchtext\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mdatasets\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m IMDB\n\u001b[32m      6\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorchtext\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mdatasets\u001b[39;00m\u001b[34;01m.\u001b[39;00m\u001b[34;01mimdb\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m NUM_LINES\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/site-packages/torchtext/__init__.py:18\u001b[39m\n\u001b[32m     15\u001b[39m     _WARN = \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[32m     17\u001b[39m \u001b[38;5;66;03m# the following import has to happen first in order to load the torchtext C++ library\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m18\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorchtext\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m _extension  \u001b[38;5;66;03m# noqa: F401\u001b[39;00m\n\u001b[32m     20\u001b[39m _TEXT_BUCKET = \u001b[33m\"\u001b[39m\u001b[33mhttps://download.pytorch.org/models/text/\u001b[39m\u001b[33m\"\u001b[39m\n\u001b[32m     22\u001b[39m _CACHE_DIR = os.path.expanduser(os.path.join(_get_torch_home(), \u001b[33m\"\u001b[39m\u001b[33mtext\u001b[39m\u001b[33m\"\u001b[39m))\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/site-packages/torchtext/_extension.py:64\u001b[39m\n\u001b[32m     59\u001b[39m     \u001b[38;5;66;03m# This import is for initializing the methods registered via PyBind11\u001b[39;00m\n\u001b[32m     60\u001b[39m     \u001b[38;5;66;03m# This has to happen after the base library is loaded\u001b[39;00m\n\u001b[32m     61\u001b[39m     \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorchtext\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m _torchtext  \u001b[38;5;66;03m# noqa\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m64\u001b[39m \u001b[43m_init_extension\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/site-packages/torchtext/_extension.py:58\u001b[39m, in \u001b[36m_init_extension\u001b[39m\u001b[34m()\u001b[39m\n\u001b[32m     55\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m _mod_utils.is_module_available(\u001b[33m\"\u001b[39m\u001b[33mtorchtext._torchtext\u001b[39m\u001b[33m\"\u001b[39m):\n\u001b[32m     56\u001b[39m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m(\u001b[33m\"\u001b[39m\u001b[33mtorchtext C++ Extension is not found.\u001b[39m\u001b[33m\"\u001b[39m)\n\u001b[32m---> \u001b[39m\u001b[32m58\u001b[39m \u001b[43m_load_lib\u001b[49m\u001b[43m(\u001b[49m\u001b[33;43m\"\u001b[39;49m\u001b[33;43mlibtorchtext\u001b[39;49m\u001b[33;43m\"\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[32m     59\u001b[39m \u001b[38;5;66;03m# This import is for initializing the methods registered via PyBind11\u001b[39;00m\n\u001b[32m     60\u001b[39m \u001b[38;5;66;03m# This has to happen after the base library is loaded\u001b[39;00m\n\u001b[32m     61\u001b[39m \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34;01mtorchtext\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;28;01mimport\u001b[39;00m _torchtext\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/site-packages/torchtext/_extension.py:50\u001b[39m, in \u001b[36m_load_lib\u001b[39m\u001b[34m(lib)\u001b[39m\n\u001b[32m     48\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m path.exists():\n\u001b[32m     49\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mFalse\u001b[39;00m\n\u001b[32m---> \u001b[39m\u001b[32m50\u001b[39m \u001b[43mtorch\u001b[49m\u001b[43m.\u001b[49m\u001b[43mops\u001b[49m\u001b[43m.\u001b[49m\u001b[43mload_library\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     51\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;01mTrue\u001b[39;00m\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/site-packages/torch/_ops.py:1392\u001b[39m, in \u001b[36m_Ops.load_library\u001b[39m\u001b[34m(self, path)\u001b[39m\n\u001b[32m   1387\u001b[39m path = _utils_internal.resolve_library_path(path)\n\u001b[32m   1388\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m dl_open_guard():\n\u001b[32m   1389\u001b[39m     \u001b[38;5;66;03m# Import the shared library into the process, thus running its\u001b[39;00m\n\u001b[32m   1390\u001b[39m     \u001b[38;5;66;03m# static (global) initialization code in order to register custom\u001b[39;00m\n\u001b[32m   1391\u001b[39m     \u001b[38;5;66;03m# operators with the JIT.\u001b[39;00m\n\u001b[32m-> \u001b[39m\u001b[32m1392\u001b[39m     \u001b[43mctypes\u001b[49m\u001b[43m.\u001b[49m\u001b[43mCDLL\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpath\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   1393\u001b[39m \u001b[38;5;28mself\u001b[39m.loaded_libraries.add(path)\n",
      "\u001b[36mFile \u001b[39m\u001b[32m~/.conda/envs/dl/lib/python3.11/ctypes/__init__.py:376\u001b[39m, in \u001b[36mCDLL.__init__\u001b[39m\u001b[34m(self, name, mode, handle, use_errno, use_last_error, winmode)\u001b[39m\n\u001b[32m    373\u001b[39m \u001b[38;5;28mself\u001b[39m._FuncPtr = _FuncPtr\n\u001b[32m    375\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m handle \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[32m--> \u001b[39m\u001b[32m376\u001b[39m     \u001b[38;5;28mself\u001b[39m._handle = \u001b[43m_dlopen\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_name\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mmode\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    377\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m    378\u001b[39m     \u001b[38;5;28mself\u001b[39m._handle = handle\n",
      "\u001b[31mOSError\u001b[39m: /home/jcheng/.conda/envs/dl/lib/python3.11/site-packages/torchtext/lib/libtorchtext.so: undefined symbol: _ZN5torch6detail10class_baseC2ERKSsS3_SsRKSt9type_infoS6_"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torchtext\n",
    "from torchtext.datasets import IMDB\n",
    "from torchtext.datasets.imdb import NUM_LINES\n",
    "from torchtext.data import get_tokenizer\n",
    "from torchtext.vocab import build_vocab_from_iterator\n",
    "from torchtext.data.functional import to_map_style_dataset\n",
    "\n",
    "import sys\n",
    "import os\n",
    "import logging\n",
    "logging.basicConfig(\n",
    "    level=logging.WARN,\n",
    "    stream=sys.stdout,\n",
    "    format=\"%(asctime)s (%(module)s:%(lineno)d) %(levelname)s: %(message)s\",\n",
    ")\n",
    "\n",
    "VOCAB_SIZE = 15000\n",
    "# 第一期： 编写GCNN模型代码\n",
    "class GCNN(nn.Module):\n",
    "    def __init__(self, vocab_size=VOCAB_SIZE, embedding_dim=64, num_class=2):\n",
    "        super(GCNN, self).__init__()\n",
    "\n",
    "        self.embedding_table = nn.Embedding(vocab_size, embedding_dim)\n",
    "        nn.init.xavier_uniform_(self.embedding_table.weight)\n",
    "\n",
    "        self.conv_A_1 = nn.Conv1d(embedding_dim, 64, 15, stride=7)\n",
    "        self.conv_B_1 = nn.Conv1d(embedding_dim, 64, 15, stride=7)\n",
    "\n",
    "        self.conv_A_2 = nn.Conv1d(64, 64, 15, stride=7)\n",
    "        self.conv_B_2 = nn.Conv1d(64, 64, 15, stride=7)\n",
    "\n",
    "        self.output_linear1 = nn.Linear(64, 128)\n",
    "        self.output_linear2 = nn.Linear(128, num_class)\n",
    "\n",
    "    def forward(self, word_index):\n",
    "        # 定义GCN网络的算子操作流程，基于句子单词ID输入得到分类logits输出\n",
    "\n",
    "        # 1. 通过word_index得到word_embedding\n",
    "        # word_index shape:[bs, max_seq_len]\n",
    "        word_embedding = self.embedding_table(word_index) #[bs, max_seq_len, embedding_dim]\n",
    "\n",
    "        # 2. 编写第一层1D门卷积模块\n",
    "        word_embedding = word_embedding.transpose(1, 2) #[bs, embedding_dim, max_seq_len]\n",
    "        A = self.conv_A_1(word_embedding)\n",
    "        B = self.conv_B_1(word_embedding)\n",
    "        H = A * torch.sigmoid(B) #[bs, 64, max_seq_len]\n",
    "\n",
    "        A = self.conv_A_2(H)\n",
    "        B = self.conv_B_2(H)\n",
    "        H = A * torch.sigmoid(B) #[bs, 64, max_seq_len]\n",
    "\n",
    "        # 3. 池化并经过全连接层\n",
    "        pool_output = torch.mean(H, dim=-1) #平均池化，得到[bs, 64]\n",
    "        linear1_output = self.output_linear1(pool_output)\n",
    "        logits = self.output_linear2(linear1_output) #[bs, 2]\n",
    "\n",
    "        return logits\n",
    "\n",
    "\n",
    "class TextClassificationModel(nn.Module):\n",
    "    \"\"\" 简单版embeddingbag+DNN模型 \"\"\"\n",
    "\n",
    "    def __init__(self, vocab_size=VOCAB_SIZE, embed_dim=64, num_class=2):\n",
    "        super(TextClassificationModel, self).__init__()\n",
    "        self.embedding = nn.EmbeddingBag(vocab_size, embed_dim, sparse=False)\n",
    "        self.fc = nn.Linear(embed_dim, num_class)\n",
    "\n",
    "    def forward(self, token_index):\n",
    "        embedded = self.embedding(token_index) # shape: [bs, embedding_dim]\n",
    "        return self.fc(embedded)\n",
    "\n",
    "\n",
    "\n",
    "# step2 构建IMDB DataLoader\n",
    "\n",
    "BATCH_SIZE = 64\n",
    "\n",
    "def yield_tokens(train_data_iter, tokenizer):\n",
    "    for i, sample in enumerate(train_data_iter):\n",
    "        label, comment = sample\n",
    "        yield tokenizer(comment)\n",
    "\n",
    "train_data_iter = IMDB(root='./imdb_dataset', split='train') # Dataset类型的对象\n",
    "tokenizer = get_tokenizer(\"basic_english\")\n",
    "vocab = build_vocab_from_iterator(yield_tokens(train_data_iter, tokenizer), min_freq=20, specials=[\"<unk>\"])\n",
    "vocab.set_default_index(0)\n",
    "print(f\"单词表大小: {len(vocab)}\")\n",
    "\n",
    "def collate_fn(batch):\n",
    "    \"\"\" 对DataLoader所生成的mini-batch进行后处理 \"\"\"\n",
    "    target = []\n",
    "    token_index = []\n",
    "    max_length = 0\n",
    "    for i, (label, comment) in enumerate(batch):\n",
    "        tokens = tokenizer(comment)\n",
    "\n",
    "        token_index.append(vocab(tokens))\n",
    "        if len(tokens) > max_length:\n",
    "            max_length = len(tokens)\n",
    "\n",
    "        if label == \"pos\":\n",
    "            target.append(0)\n",
    "        else:\n",
    "            target.append(1)\n",
    "\n",
    "    token_index = [index + [0]*(max_length-len(index)) for index in token_index]\n",
    "    return (torch.tensor(target).to(torch.int64), torch.tensor(token_index).to(torch.int32))\n",
    "\n",
    "\n",
    "# step3 编写训练代码\n",
    "def train(train_data_loader, eval_data_loader, model, optimizer, num_epoch, log_step_interval, save_step_interval, eval_step_interval, save_path, resume=\"\"):\n",
    "    \"\"\" 此处data_loader是map-style dataset \"\"\"\n",
    "    start_epoch = 0\n",
    "    start_step = 0\n",
    "    if resume != \"\":\n",
    "        #  加载之前训过的模型的参数文件\n",
    "        logging.warning(f\"loading from {resume}\")\n",
    "        checkpoint = torch.load(resume)\n",
    "        model.load_state_dict(checkpoint['model_state_dict'])\n",
    "        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n",
    "        start_epoch = checkpoint['epoch']\n",
    "        start_step = checkpoint['step']\n",
    "\n",
    "    for epoch_index in range(start_epoch, num_epoch):\n",
    "        ema_loss = 0.\n",
    "        num_batches = len(train_data_loader)\n",
    "\n",
    "        for batch_index, (target, token_index) in enumerate(train_data_loader):\n",
    "            optimizer.zero_grad()\n",
    "            step = num_batches*(epoch_index) + batch_index + 1\n",
    "            logits = model(token_index)\n",
    "            bce_loss = F.binary_cross_entropy(torch.sigmoid(logits), F.one_hot(target, num_classes=2).to(torch.float32))\n",
    "            ema_loss = 0.9*ema_loss + 0.1*bce_loss\n",
    "            bce_loss.backward()\n",
    "            nn.utils.clip_grad_norm_(model.parameters(), 0.1)\n",
    "            optimizer.step()\n",
    "\n",
    "            if step % log_step_interval == 0:\n",
    "                logging.warning(f\"epoch_index: {epoch_index}, batch_index: {batch_index}, ema_loss: {ema_loss.item()}\")\n",
    "\n",
    "            if step % save_step_interval == 0:\n",
    "                os.makedirs(save_path, exist_ok=True)\n",
    "                save_file = os.path.join(save_path, f\"step_{step}.pt\")\n",
    "                torch.save({\n",
    "                    'epoch': epoch_index,\n",
    "                    'step': step,\n",
    "                    'model_state_dict': model.state_dict(),\n",
    "                    'optimizer_state_dict': optimizer.state_dict(),\n",
    "                    'loss': bce_loss,\n",
    "                }, save_file)\n",
    "                logging.warning(f\"checkpoint has been saved in {save_file}\")\n",
    "\n",
    "            if step % eval_step_interval == 0:\n",
    "                logging.warning(\"start to do evaluation...\")\n",
    "                model.eval()\n",
    "                ema_eval_loss = 0\n",
    "                total_acc_account = 0\n",
    "                total_account = 0\n",
    "                for eval_batch_index, (eval_target, eval_token_index) in enumerate(eval_data_loader):\n",
    "                    total_account += eval_target.shape[0]\n",
    "                    eval_logits = model(eval_token_index)\n",
    "                    total_acc_account += (torch.argmax(eval_logits, dim=-1) == eval_target).sum().item()\n",
    "                    eval_bce_loss = F.binary_cross_entropy(torch.sigmoid(eval_logits), F.one_hot(eval_target, num_classes=2).to(torch.float32))\n",
    "                    ema_eval_loss = 0.9*ema_eval_loss + 0.1*eval_bce_loss\n",
    "                acc = total_acc_account/total_account\n",
    "\n",
    "                logging.warning(f\"eval_ema_loss: {ema_eval_loss.item()}, eval_acc: {acc.item()}\")\n",
    "                model.train()\n",
    "\n",
    "# step4 测试代码\n",
    "if __name__ == \"__main__\":\n",
    "    model = GCNN()\n",
    "    #  model = TextClassificationModel()\n",
    "    print(\"模型总参数:\", sum(p.numel() for p in model.parameters()))\n",
    "    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n",
    "\n",
    "    train_data_iter = IMDB(root='.data', split='train') # Dataset类型的对象\n",
    "    train_data_loader = torch.utils.data.DataLoader(to_map_style_dataset(train_data_iter), batch_size=BATCH_SIZE, collate_fn=collate_fn, shuffle=True)\n",
    "\n",
    "    eval_data_iter = IMDB(root='.data', split='test') # Dataset类型的对象\n",
    "    eval_data_loader = torch.utils.data.DataLoader(to_map_style_dataset(eval_data_iter), batch_size=8, collate_fn=collate_fn)\n",
    "    resume = \"\"\n",
    "\n",
    "    train(train_data_loader, eval_data_loader, model, optimizer, num_epoch=10, log_step_interval=20, save_step_interval=500, eval_step_interval=300, save_path=\"./logs_imdb_text_classification\", resume=resume)\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dl",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
